"""Implements tools for drawing 2d plots.

This module implements tools for drawing contour plots 
of objective functions and iterates on top of contour 
lines.
"""

import math
import numpy
from native import *
from plotutils import *
from testproblems import *
from matplotlib.pyplot import *
from numpy import arange,logspace,meshgrid

def plot_bound_constraints(C):
	plot([C.L[0], C.U[0]], [C.L[1], C.L[1]], color='black', linestyle='dashed')
	plot([C.L[0], C.U[0]], [C.U[1], C.U[1]], color='black', linestyle='dashed')
	plot([C.L[0], C.L[0]], [C.L[1], C.U[1]], color='black', linestyle='dashed')
	plot([C.U[0], C.U[0]], [C.L[1], C.U[1]], color='black', linestyle='dashed')

def plot_contours(f, x1, x2, y1, y2, z1, z2,
                  z_logscale, new_figure=True):
	"""Plots contour lines of the given objective function. The 
plotting range is [x1,x2]x[y1,y2]x[z1,z2], and logarithmic 
scaling of z-axis is specified with the boolean argument 
z_logscale. If new_figure is set to true, this function 
opens a new window for the plot.
"""
	X,Y,Z = tabulate_function(f, 300, (x1, x2), (y1, y2))
  
	if new_figure == True:
		fig = figure()
  
	if z_logscale == True:
		V = logspace(math.log10(z1), math.log10(z2), 15)
	else:
		V = arange(z1, z2, (z2 - z1) / 20);
	
	contour(X, Y, Z, V, colors='k', linewidths=0.25)
  
	xlim(x1, x2)
	ylim(y1, y2)
  
	if new_figure == True:
		show()

def plot_iterates_on_contours(results, x1, x2, y1, y2, z1, z2,
                              z_logscale=True, new_figure=True):
	"""Given a results object generated by calling pyotk_driver.minimize(),
plot the iteration path on contour lines of the objective function. 
The plotting range is [x1,x2]x[y1,y2]x[z1,z2], and logarithmic 
scaling of z-axis is specified with the boolean argument 
z_logscale. If new_figure is set to true, this function opens 
a new window for the plot. This function requires that 
the objective function is a function of two variables, otherwise 
a ValueError is raised."""
	if new_figure == True:
		figure()
	
	if results.n != 2:
		raise ValueError('cannot plot an objective function of more than two variables')
	
	plot_contours(results.input.objfunc, x1, x2, y1, y2, z1, z2,
                z_logscale=z_logscale, new_figure=False)
	
	PX = []
	PY = []
	
	if results.input.solver_info.m == 1:
		for it in results.states:
			PX.append(it.x[0])
			PY.append(it.x[1])
			
		plot(PX, PY, '-o', color='black',
         markerfacecolor='gray', markeredgecolor='black',
         linewidth=2, markersize=6)
	else:
		for s in results.states:
			PX = []
			PY = []
			for vi in range(0, 4):
				PX.append(s.X[0, vi % 3])
				PY.append(s.X[1, vi % 3])
		
			plot(PX, PY, '-o', color='black',
           markerfacecolor='gray', markeredgecolor='black',
           linewidth=2, markersize=6)
	
	if results.input.has_constraints and type(results.input.C) == native.BoundConstraints:
		plot_bound_constraints(results.input.C)
	
	xlim(x1, x2)
	ylim(y1, y2)
	
	if new_figure == True:
		show()

def main():
	f = Beale()
	x1 = f.plot_spec.x_range[0]
	x2 = f.plot_spec.x_range[1]
	y1 = f.plot_spec.y_range[0]
	y2 = f.plot_spec.y_range[1]
	z1 = f.plot_spec.z_range[0]
	z2 = f.plot_spec.z_range[1]
	z_logscale = f.plot_spec.z_logscale
	#plot_contours(f.otk_instance, x1, x2, y1, y2, z1, z2, z_logscale)
	
	s = GSLfdfsolver('vector_bfgs2')
	results = minimize(s, DefaultSolverSetup(), f.otk_instance, f.stopcrit, f.x0, NoConstraints(), 0, False)
	plot_iterates_on_contours(results, x1, x2, y1, y2, z1, z2)

	f = Trigonometric(n=2)
	x1 = -1
	x2 = 1
	y1 = -0.5
	y2 = 1
	z1 = f.plot_spec.z_range[0]
	z2 = f.plot_spec.z_range[1]
	z_logscale = f.plot_spec.z_logscale
	
	s = LBFGSB()
	C = BoundConstraints(2)
	C.types[0] = BoundType.both
	C.types[1] = BoundType.both
	C.L = [-0.25, -0.25]
	C.U = [0.25, 0.25]
	x0 = (-0.5, 0.5)
	results = minimize(s, DefaultSolverSetup(), f.otk_instance, f.stopcrit, x0, C, 0, False)
	plot_iterates_on_contours(results, x1, x2, y1, y2, z1, z2)

	C.L = [0.0, 0.4]
	C.U = [0.6, 0.75]
	results = minimize(s, DefaultSolverSetup(), f.otk_instance, f.stopcrit, x0, C, 0, False)
	plot_iterates_on_contours(results, x1, x2, y1, y2, z1, z2)
	
	s = GSLfsolver('nmsimplex')
	f = PenaltyFunctionII(n=2)
	x1 = f.plot_spec.x_range[0]
	x2 = f.plot_spec.x_range[1]
	y1 = f.plot_spec.y_range[0]
	y2 = f.plot_spec.y_range[1]
	z1 = f.plot_spec.z_range[0]
	z2 = f.plot_spec.z_range[1]
	z_logscale = f.plot_spec.z_logscale
	x0 = (-0.75, 0.5)
	results = minimize(s, DefaultSolverSetup(), f.otk_instance, f.stopcrit, x0, C, 0, False)
	plot_iterates_on_contours(results, x1, x2, y1, y2, z1, z2)
	
if __name__ == "__main__":
	main()
